#include <cstdio>
#include <cstdlib>
#include <vector>
#include "utils.h"
#include <cuda_runtime.h>
#include <random>

#include "check_point/reduce_ck0.cpp"

int main(int argc, char *argv[]) {
    cudaSetDevice(0);
    std::vector<Case>  cases = {};

    #ifdef PERF

    float m_start = 100000;
    int loop = 1000;
    
    for(int m_id = m_start; m_id <= 10*m_start; m_id+=m_start){
        cases.push_back({m_id, loop});
    }

    #endif

    #ifdef EVAL
    const int m = atoi(argv[1]), loop = atoi(argv[2]);
    const std::string src_gt_path = argv[3], dst_gt_path = argv[4], dst_my_path = argv[5];
    cases.push_back({m, loop});
    #endif

    #ifdef PROFILE
    cases.push_back({262133, 1});
    #endif

    for(size_t case_id = 0; case_id < cases.size(); ++case_id){
        #ifdef PERF
        if(case_id > 0){
            cases[case_id].loop = max(1, int(1000.0f/cases[case_id-1].avg_latency))*3;
        }
        #endif

        Case c = cases[case_id];
        const int m = c.m, loop = c.loop;
        cudaStream_t stream = NULL;
        CUDA_CHECK(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));
        float *d_src = nullptr;
        float *d_tmp_dst = nullptr;
        float *d_dst = nullptr;
        CUDA_CHECK(cudaMalloc(reinterpret_cast<void **>(&d_src), sizeof(float)*m));
        CUDA_CHECK(cudaMalloc(reinterpret_cast<void **>(&d_dst), sizeof(float)*1));
        const int deal_num_per_bloc = 1024;
        int num_block = DIV_UP(m, deal_num_per_bloc);
        std::vector<int> its;
        its.push_back(m);
        if(num_block > 1){
            CUDA_CHECK(cudaMalloc(reinterpret_cast<void **>(&d_tmp_dst), sizeof(float)*num_block));
            for(int i = num_block; i > 1; i = DIV_UP(i, deal_num_per_bloc)){
                its.push_back(i);
            }
        }
        its.push_back(1);

        #ifdef EVAL
        loadbin(d_src, m*sizeof(float), src_gt_path);
        #endif
        {
            int start = now();
            for(int i = 0; i < loop; ++i){
                if(its.size() == 2){
                    reduce(its[0], d_src, d_dst, stream);
                }else{
                    for(int i = 0; i < its.size(); ++i){
                        if(its[i] > 1 && its[i+1] != 1){
                            if(i == 0){
                                reduce(its[i], d_src, d_tmp_dst, stream);
                            }else{
                                reduce(its[i], d_tmp_dst, d_tmp_dst, stream);
                            }
                        }else if(its[i] > 1 && its[i+1] == 1){
                            reduce(its[i], d_tmp_dst, d_dst, stream);                        
                        }
                    }
                }
            }
            cudaStreamSynchronize(stream);
            int end = now();
            cases[case_id].avg_latency =float(end-start)/float(loop);

            #ifdef EVAL
            dumpbin(d_dst, 1*sizeof(float), dst_my_path);
            #endif
        }
        #ifdef PERF
        {
            Case c = cases[case_id];
            const int m = c.m, loop = c.loop;
            float avg_latency = c.avg_latency;
            printf("m %d loop %d avg_latency %0.4f ms  \n", m, loop, avg_latency);
        }
        #endif
        CUDA_CHECK(cudaFree(d_src));
        CUDA_CHECK(cudaFree(d_dst));
        CUDA_CHECK(cudaStreamDestroy(stream));
    }
    return 0;
}